import json
import os
import time
import shutil
# 获取当前日期
from datetime import datetime
# 修改导入方式，使用相对导入或者绝对导入
import sys
import os as sys_os
# 获取当前脚本所在目录
current_dir = sys_os.path.dirname(sys_os.path.abspath(__file__))
# 将当前目录添加到系统路径中
if current_dir not in sys.path:
    sys.path.append(current_dir)

# 监测目录路径
monitor_dir = r"D:\my_dataset_from_server"
images_directory = rf"{monitor_dir}\images"
labels_directory = rf"{monitor_dir}\labels"
# 目标目录路径
target_dir = r"D:\myGitee\yolo-universal-training\my_dataset"
# 运用pt的远程文件夹
network_path = r"\\192.168.11.123\models"
# 图片数量阈值
# image_threshold = 80
# 训练任务参数
train_task_input = "检测"
train_model_input = r"best.pt"
device_input = "cuda"
data_input = "my_dataset"
# epochs_input = 200
batch_size_input = 16
imgsz_input = 640
include_test = True
test_size = 0.2
val_size = 0.2

patience_input = 50
optimizer_input = "Adam"
lr0_input = 0.001
lrf_input = 0.01
momentum_input = 0.937
weight_decay_input = 0.0005
warmup_epochs_input = 3.0
warmup_momentum_input = 0.8
warmup_bias_lr_input = 0.1
box_input = 7.5
cls_input = 0.5
dfl_input = 1.5
pose_input = 12.0
kobj_input = 2.0
label_smoothing_input = 0.1
nbs_input = 64
overlap_mask_input = True
mask_ratio_input = 4
dropout_input = 0.2

# 新增的高级选项
augment_input = True
mixup_input = 0.2
copy_paste_input = 0.2
mosaic_input = 1.0

# 新增的 AMP 选项
amp_input = True

# 备份目录路径
backup_images_dir = r"D:\backup\images"
backup_labels_dir = r"D:\backup\labels"


def count_images_in_directory(directory):
    """
    统计目录下图片文件的数量
    """
    if not os.path.exists(directory):
        return 0

    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}
    count = 0

    # 使用os.walk递归遍历所有子目录
    for root, dirs, files in os.walk(directory):
        for filename in files:
            if os.path.splitext(filename)[1].lower() in image_extensions:
                count += 1

    return count


def copy_directory_contents(src, dst):
    """
    复制源目录内容到目标目录，完成后删除已复制的源目录内容，返回已经复制的文件列表
    """
    if not os.path.exists(dst):
        os.makedirs(dst)

    copied_files = []  # 记录已复制的文件
    for item in os.listdir(src):
        src_path = os.path.join(src, item)
        dst_path = os.path.join(dst, item)

        if os.path.isdir(src_path):
            # 如果是目录，则遍历目录中的文件
            for root, dirs, files in os.walk(src_path):
                # 计算相对路径
                rel_path = os.path.relpath(root, src)
                dst_root = os.path.join(dst, rel_path)
                
                # 创建目标目录
                if not os.path.exists(dst_root):
                    os.makedirs(dst_root)
                
                # 复制文件
                for file in files:
                    src_file = os.path.join(root, file)
                    dst_file = os.path.join(dst_root, file)
                    shutil.copy2(src_file, dst_file)
                    # 记录相对路径和文件名
                    copied_files.append(os.path.join(rel_path, src_file))
        else:
            # 复制文件
            shutil.copy2(src_path, dst_path)
            copied_files.append(item)  # 记录已复制的文件名

    return copied_files  # 返回已复制的文件列表


def delete_copied_contents(src, copied_items):
    """
    只删除已复制的内容，保留新添加的内容
    """
    print(f"删除{src}的已复制内容,共计{len(copied_items)}个文件...")
    for item in copied_items:
        src_path = os.path.join(src, item)
        if os.path.exists(src_path):
            if os.path.isfile(src_path):
                os.remove(src_path)


def move_datasetfiles_to_backup(target_dir):
    """
    将目标目录下的images文件夹中的jpg文件剪切到备份目录，
    将labels文件夹中的txt文件剪切到备份目录
    """
    print(f"备份数据集文件到{target_dir}...")
    # 确保备份目录存在
    if not os.path.exists(backup_images_dir):
        os.makedirs(backup_images_dir)
    if not os.path.exists(backup_labels_dir):
        os.makedirs(backup_labels_dir)
    print("1.移动数据集文件...")
    # 移动images文件夹中的图片文件（不含子文件夹）
    images_dir = os.path.join(target_dir, "images")
    if os.path.exists(images_dir):
        for file in os.listdir(images_dir):
            # 支持更多图像格式
            if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp')):
                src_path = os.path.join(images_dir, file)
                if os.path.isfile(src_path):  # 确保只移动文件，不移动文件夹
                    dst_path = os.path.join(backup_images_dir, file)
                    try:
                        shutil.move(src_path, dst_path)
                    except Exception as e:
                        print(f"移动图像文件 {src_path} 时出错: {e}")
    print("2.移动标签文件...")
    # 移动labels文件夹中的txt文件（不含子文件夹）
    labels_dir = os.path.join(target_dir, "labels")
    if os.path.exists(labels_dir):
        for file in os.listdir(labels_dir):
            if file.lower().endswith('.txt'):
                src_path = os.path.join(labels_dir, file)
                if os.path.isfile(src_path):  # 确保只移动文件，不移动文件夹
                    dst_path = os.path.join(backup_labels_dir, file)
                    try:
                        shutil.move(src_path, dst_path)
                    except Exception as e:
                        print(f"移动标签文件 {src_path} 时出错: {e}")


def copy_best_model_to_network(save_dir):
    """
    将app中的save_dir中最新的best.pt文件复制到网络路径，并重命名
    """
    try:
        print("开始复制best.pt文件...")
        # 检查save_dir是否存在
        if not os.path.exists(save_dir):
            print(f"保存目录 {save_dir} 不存在")
            return

        # 查找最新的best.pt文件
        best_model_path = os.path.join(save_dir, "weights", "best.pt")
        if not os.path.exists(best_model_path):
            print(f"best.pt文件不存在于 {best_model_path}")
            return

        # current_date = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 构建目标文件路径
        if not os.path.exists(network_path):
            print(f"网络路径 {network_path} 不存在")
            return

        # target_filename = f"best_{current_date}.pt"
        target_path = os.path.join(network_path, "best.pt")

        # 根据修改时间重命名network_path中已有的best.pt文件
        existing_best_path = os.path.join(network_path, "best.pt")
        if os.path.exists(existing_best_path):
            # 获取文件的修改时间
            mod_time = os.path.getmtime(existing_best_path)
            mod_time_str = datetime.fromtimestamp(mod_time).strftime("%Y%m%d_%H%M%S")
            # 重命名为best_时间.pt格式
            backup_name = f"best_{mod_time_str}.pt"
            backup_path = os.path.join(network_path, backup_name)
            print(f"备份{network_path}至{backup_path}")
            os.rename(existing_best_path, backup_path)
            print(f"原best.pt已根据修改时间重命名为 {backup_path}")
        # 复制文件到网络路径
        shutil.copy2(best_model_path, target_path)
        print(f"成功将best.pt复制到{network_path} {target_path}")

        # 复制一份到train_model_input路径下
        original_best_path = r"models\detection\best.pt"
        train_model_dir = os.path.dirname(original_best_path)
        if not os.path.exists(train_model_dir):
            os.makedirs(train_model_dir)

        # 备份原有的best.pt文件
        if os.path.exists(original_best_path):
            # 获取文件的修改时间
            mod_time = os.path.getmtime(original_best_path)
            mod_time_str = datetime.fromtimestamp(mod_time).strftime("%Y%m%d_%H%M%S")
            # 根据修改时间重命名备份文件
            backup_name = f"best_backup_{mod_time_str}.pt"
            backup_path = os.path.join(train_model_dir, backup_name)
            print(f"备份{original_best_path}至{backup_path}")
            shutil.copy2(original_best_path, backup_path)
            print(f"原best.pt已备份为 {backup_path}")
        # 复制新的best.pt到train_model_input路径
        shutil.copy2(best_model_path, original_best_path)
        print(f"成功将best.pt复制到{backup_path} {original_best_path}")
    except Exception as e:
        print(f"复制best.pt文件时出错: {e}")

# 获取目录中的所有文件名（不含扩展名）
def get_filenames_without_extension(directory, extension):
    filenames = []
    for file in os.listdir(directory):
        if file.endswith(extension):
            filenames.append(os.path.splitext(file)[0])
    return set(filenames)

# 删除不匹配的文件
def delete_unmatched_files(images_dir, labels_dir):
    # 获取jpg文件名（不含扩展名）
    jpg_files = get_filenames_without_extension(images_dir, '.jpg')
    # 获取txt文件名（不含扩展名）
    txt_files = get_filenames_without_extension(labels_dir, '.txt')
    
    # 找出只在images目录中存在的文件
    unmatched_images = jpg_files - txt_files
    # 找出只在labels目录中存在的文件
    unmatched_labels = txt_files - jpg_files
    
    # 删除images目录中没有对应txt的jpg文件
    for filename in unmatched_images:
        file_path = os.path.join(images_dir, filename + '.jpg')
        if os.path.exists(file_path):
            print(f"删除: {file_path}")
            os.remove(file_path)
    
    # 删除labels目录中没有对应jpg的txt文件
    for filename in unmatched_labels:
        file_path = os.path.join(labels_dir, filename + '.txt')
        if os.path.exists(file_path):
            print(f"删除: {file_path}")
            os.remove(file_path)
def main():
    # 实时监测目录下的图片数量
    print(f"开始监测目录 {monitor_dir}...")
    while True:
        with open("config/settings.json", "r") as f:
            config = json.load(f)
            image_threshold = config["image_threshold"]
            epochs_input = config["epochs_input"]
        image_count = count_images_in_directory(monitor_dir)
        original_best_path = r"models\detection\best.pt"
        mod_time = os.path.getmtime(original_best_path)
        mod_time_str = datetime.fromtimestamp(mod_time).strftime("%Y年%m月%d日 %H时%M分%S秒")
        print(
            f"目录 {monitor_dir} 下检测到 {image_count} 张图片,阈值 {image_threshold}张, 轮数 {epochs_input} ，使用best.pt文件为{mod_time_str}",
            end='\r')
        if image_count >= image_threshold:
            print(f"\n图片数量已达到 {image_threshold} 张，开始复制文件...")            
            delete_unmatched_files(images_directory, labels_directory)
            # 确保目标目录存在
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)

            # 复制所有文件到目标目录，不清空目标目录
            #
            if os.path.exists(monitor_dir):
                # 复制目录内容到目标位置，保留目标目录原有内容
                copied_items = copy_directory_contents(monitor_dir, target_dir)

            print(f"文件复制完成，开始执行训练任务...")

            # 现在可以导入app模块了
            from app import train_model
            resultstr, save_dir = train_model(
                data_input,
                epochs_input,
                train_task_input,
                train_model_input,
                device_input,
                include_test,
                test_size,
                val_size,
                batch_size_input,
                imgsz_input,
                patience_input,
                optimizer_input,
                lr0_input,
                lrf_input,
                momentum_input,
                weight_decay_input,
                warmup_epochs_input,
                warmup_momentum_input,
                warmup_bias_lr_input,
                box_input,
                cls_input,
                dfl_input,
                pose_input,
                kobj_input,
                label_smoothing_input,
                nbs_input,
                overlap_mask_input,
                mask_ratio_input,
                dropout_input,
                # 新增的输入
                augment_input,
                mixup_input,
                copy_paste_input,
                mosaic_input,
                amp_input)
            print(resultstr)
            # 训练完成后最大程度释放资源
            print("训练完成，开始释放资源...")

            # 强制进行垃圾回收
            print("开始垃圾回收")
            import gc
            gc.collect()
            print("垃圾回收完成")
            # 清理CUDA缓存（如果使用了GPU）
            try:
                print("清理CUDA缓存")
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
            except ImportError:
                pass
            print("CUDA缓存清理完成")

            # 将训练数据移动到备份目录
            move_datasetfiles_to_backup(target_dir)

            # 复制best.pt到网络路径
            copy_best_model_to_network(save_dir)

            # 删除已复制的源文件夹内容
            delete_copied_contents(monitor_dir, copied_items)
            # 重启自身程序以完全释放资源
            print("资源释放完成，重启程序...")
            import sys
            import subprocess
            # 使用当前Python解释器重新运行当前脚本
            subprocess.Popen([sys.executable] + sys.argv)
            # 退出当前进程
            sys.exit(0)
        else:
            time.sleep(5)  # 每5秒检查一次


if __name__ == "__main__":
    main()
    # copy_best_model_to_network()
